package com.gjy.console.ssh;

import lombok.Getter;
import lombok.Setter;
import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.keyverifier.AcceptAllServerKeyVerifier;
import org.apache.sshd.client.session.ClientSession;
import org.apache.sshd.common.keyprovider.FileKeyPairProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.file.Paths;

/**
 * @author gjy
 * @version 1.0
 * @since 2025-09-04 18:17:07
 */
@Getter
@Setter
public class SshSession {
    private static final Logger log = LoggerFactory.getLogger(SshSession.class);

    public static final Integer UP = 1;
    public static final Integer PK = 2;

    private String host;
    private int port;
    private String username;
    private String password;
    private int timeout;
    private String privateKey;

    public SshSession(String host, int port, String username, int timeout, String privateKey) {
        this.host = host;
        this.port = port;
        this.username = username;
        this.timeout = timeout;
        this.privateKey = privateKey;
    }

    public SshSession() {
    }

    public SshSession(String host, int port, String username, String password, int timeout) {
        this.host = host;
        this.port = port;
        this.username = username;
        this.password = password;
        this.timeout = timeout;
    }

    private SshClient client;
    private ClientSession session;

    /**
     * 建立SSH连接
     */
    public ClientSession connectWithUp() throws IOException {
        if (session != null && session.isOpen()) {
            return session;
        }
        client = SshClient.setUpDefaultClient();
        client.start();
        session = client.connect(username, host, port)
                .verify(timeout)
                .getSession();

        session.addPasswordIdentity(password);
        session.auth().verify(timeout);
        log.info("SSH 用户密码认证成功");
        return session;
    }

    public ClientSession connectWithPk() throws IOException {
        if (session != null && session.isOpen()) {
            return session;
        }
        client = SshClient.setUpDefaultClient();
        // 生产环境请替换为已知主机密钥验证器
        client.setServerKeyVerifier(AcceptAllServerKeyVerifier.INSTANCE);
        client.start();
        session = client.connect(username, host, port).verify().getSession();
        FileKeyPairProvider keys = new FileKeyPairProvider(Paths.get(privateKey));
        session.setKeyIdentityProvider(keys);
        session.auth().verify();
        log.info("SSH 公钥认证成功");
        return session;
    }

    /**
     * 断开SSH连接
     */
    public void disconnect() throws IOException {
        if (session != null && session.isOpen()) {
            session.close();
        }
        if (client != null && client.isStarted()) {
            client.stop();
        }
    }

    public ClientSession getSession(Integer type) throws IOException {
        if (session == null || session.isEmpty() || session.isClosed()) {
            if (type == 1) {
                return connectWithUp();
            } else if (type == 2) {
                return  connectWithPk();
            } else {
                throw new IllegalArgumentException();
            }
        }
        return session;
    }

}
